import pandas as pd
from datasets import Dataset
from transformers import BertTokenizer, BertForSequenceClassification, Trainer, TrainingArguments
from sklearn.model_selection import train_test_split
import torch
torch.cuda.empty_cache()

# 存储过程注入的关键字列表
STORED_PROCEDURE_KEYWORDS = [
    "EXEC",
    "EXECUTE",
    "sp_executesql",
    "DECLARE",
    "SET",
    "OPENROWSET",
    "BEGIN",
    "END",
    "CREATE PROCEDURE",
    "ALTER PROCEDURE",
    "DROP PROCEDURE",
    "CREATE FUNCTION",
    "ALTER FUNCTION",
    "DROP FUNCTION"
]

# 强制使用 GPU（如果可用）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"🖥️ Using device: {device}")

# 1️ 加载数据
df = pd.read_csv("datasets/SQLiV3.tsv", sep="\t")
df = df.rename(columns={"payload": "text", "label": "label"})
df = df.dropna()

# 打印原始类别分布
print("👉 原始Label分布：")
print(df['label'].value_counts())

# 2️ 标记存储过程注入
def mark_stored_procedure_injection(text, label):
    # 只对已经标记为SQL注入的样本进行处理
    if label == 1:  # 假设1表示SQL注入
        # 检查文本中是否包含存储过程注入关键字
        for keyword in STORED_PROCEDURE_KEYWORDS:
            if keyword.lower() in text.lower():
                return 2  # 标记为存储过程注入
    return label

# 应用标记函数
df['label'] = df.apply(lambda row: mark_stored_procedure_injection(row['text'], row['label']), axis=1)

# 打印处理后的类别分布
print("👉 处理后的Label分布：")
print(df['label'].value_counts())

# 3️ 分割训练/测试集
train_df, test_df = train_test_split(df, test_size=0.2, random_state=42, stratify=df["label"])
train_dataset = Dataset.from_pandas(train_df)
test_dataset = Dataset.from_pandas(test_df)

# 4️ 加载英文 BERT 模型和 tokenizer
tokenizer = BertTokenizer.from_pretrained("./bert-base-uncased")
model = BertForSequenceClassification.from_pretrained("./bert-base-uncased", num_labels=3)  # 3分类：正常SQL(0)、普通SQL注入(1)、存储过程SQL注入(2)
model.to(device)

# 5️ Tokenizer 处理
def tokenize(batch):
    return tokenizer(batch["text"], padding="max_length", truncation=True, max_length=512)

train_dataset = train_dataset.map(tokenize, batched=True)
test_dataset = test_dataset.map(tokenize, batched=True)

# 6️ 设置训练参数
training_args = TrainingArguments(
    output_dir="models/bert_sql_model",
    evaluation_strategy="epoch",
    save_strategy="epoch",
    logging_dir="logs",
    num_train_epochs=3,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    save_total_limit=1,
    load_best_model_at_end=True,
    logging_steps=10,
    report_to="none",  # 禁止 wandb
)

# 7️ 定义 Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
    tokenizer=tokenizer
)

# 8️ 开始训练
print("🚀 正在开始训练 BERT 模型...")
trainer.train()

# 9️ 保存模型
model.save_pretrained("models/procedure_injection_bert")
tokenizer.save_pretrained("models/procedure_injection_bert")
print("✅ 模型和 tokenizer 已保存到 models/procedure_injection_bert")
